#include <fstream>
#include <random>
#include <thread>

#include <bthread/bthread.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

#include "hawking/indexlib/index_config.pb.h"
#include "doc/document_creator.h"
#include "index/primary_key_index/primary_key_mem_indexer.h"
#include "index/primary_key_index/primary_key_disk_indexer.h"
#include "util/file_wrapper/file_wrapper.h"

using namespace hawking::indexlib;
std::uniform_int_distribution<int32_t> dis(1, 100);
std::default_random_engine e{std::random_device{}()};

TEST(pk_mem_test, read_and_write) {
    PrimaryKeyMemIndexer pk_mem_indexer;
    pk_mem_indexer.Init();

    // add then read
    pk_mem_indexer.AddDocument(100, 0);
    pk_mem_indexer.AddDocument(101, 1);
    pk_mem_indexer.AddDocument(102, 2);
    pk_mem_indexer.AddDocument(103, 3);
    pk_mem_indexer.AddDocument(104, 4);
    pk_mem_indexer.AddDocument(105, 5);

    ReturnValue<DocId> ret = pk_mem_indexer.Lookup(100);
    EXPECT_EQ(0, ret.Value());
    ret = pk_mem_indexer.Lookup(101);
    EXPECT_EQ(1, ret.Value());
    ret = pk_mem_indexer.Lookup(102);
    EXPECT_EQ(2, ret.Value());
    ret = pk_mem_indexer.Lookup(103);
    EXPECT_EQ(3, ret.Value());
    ret = pk_mem_indexer.Lookup(104);
    EXPECT_EQ(4, ret.Value());
    ret = pk_mem_indexer.Lookup(105);
    EXPECT_EQ(5, ret.Value());
}

TEST(pk_test, read_and_write) {
    PrimaryKeyMemIndexer pk_mem_indexer;
    pk_mem_indexer.Init();
    std::shared_ptr<DeletionMapIndexer> deletion_index = std::make_shared<DeletionMapIndexer>();
    deletion_index->Init();

    // add then read
    for (uint64_t pk = 1000; pk < 102400; ++pk) {
        DocId docid = static_cast<DocId>(pk) - 1000;
        pk_mem_indexer.AddDocument(pk, docid);
    }

    std::unordered_map<DocId, DocId> old2new_map = pk_mem_indexer.GenerateOldNewDocMap(deletion_index);
    FslibWrapper::RmDir("/tmp/primary_index_test");
    FslibWrapper::Mkdir("/tmp/primary_index_test");
    pk_mem_indexer.Dump("/tmp/primary_index_test", old2new_map);

    PrimaryKeyDiskIndexer pk_disk_indexer;
    pk_disk_indexer.Load("/tmp/primary_index_test", {}, deletion_index);

    
    for (uint64_t pk = 1000; pk < 102400; ++pk) {
        ReturnValue<DocId> ret = pk_disk_indexer.Lookup(pk);
        EXPECT_EQ(ret.Code(), Status::OK);
        EXPECT_EQ(ret.Value(), static_cast<DocId>(pk) - 1000);
    }
}

TEST(pk_test, delete_dump_load) {
    PrimaryKeyMemIndexer pk_mem_indexer;
    pk_mem_indexer.Init();

    // add then read
    pk_mem_indexer.AddDocument(100, 0);
    pk_mem_indexer.AddDocument(101, 1);
    pk_mem_indexer.AddDocument(102, 2);
    pk_mem_indexer.AddDocument(103, 3);
    pk_mem_indexer.AddDocument(104, 4);
    pk_mem_indexer.AddDocument(105, 5);

    ReturnValue<DocId> ret = pk_mem_indexer.Lookup(100);
    EXPECT_EQ(0, ret.Value());
    ret = pk_mem_indexer.Lookup(101);
    EXPECT_EQ(1, ret.Value());
    ret = pk_mem_indexer.Lookup(102);
    EXPECT_EQ(2, ret.Value());
    ret = pk_mem_indexer.Lookup(103);
    EXPECT_EQ(3, ret.Value());
    ret = pk_mem_indexer.Lookup(104);
    EXPECT_EQ(4, ret.Value());
    ret = pk_mem_indexer.Lookup(105);
    EXPECT_EQ(5, ret.Value());

    // delete doc
    std::shared_ptr<DeletionMapIndexer> deletion_index = std::make_shared<DeletionMapIndexer>();
    deletion_index->Init();
    deletion_index->Delete(2);

    // dump
    std::unordered_map<DocId, DocId> old2new_map = pk_mem_indexer.GenerateOldNewDocMap(deletion_index);
    FslibWrapper::RmDir("/tmp/primary_index_test");
    FslibWrapper::Mkdir("/tmp/primary_index_test");
    pk_mem_indexer.Dump("/tmp/primary_index_test", old2new_map);

    // tomb_stone
    std::unordered_set<uint64_t> tomb_stone_pks;
    tomb_stone_pks.insert(103);
    tomb_stone_pks.insert(105);

    // load
    PrimaryKeyDiskIndexer pk_disk_indexer;
    pk_disk_indexer.Load("/tmp/primary_index_test", tomb_stone_pks, deletion_index);

    // judge
    ret = pk_disk_indexer.Lookup(102);
    EXPECT_EQ(ret.Code(), Status::NotFound);

    ret = pk_disk_indexer.Lookup(103);
    EXPECT_EQ(ret.Code(), Status::OK);
    EXPECT_EQ(true, deletion_index->IsDeleted(ret.Value()));

    ret = pk_disk_indexer.Lookup(105);
    EXPECT_EQ(ret.Code(), Status::OK);
    EXPECT_EQ(true, deletion_index->IsDeleted(ret.Value()));
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}
